# Crank-Nicolson 法求解三维Schrodinger方程
# i \hbar \pdv{u}{t} = -\frac{\hbar^2}{2m} \pdv[2]{u}{x} + V u
# \frac{u^{(k+1)}_i - u^{(k)}_i}{\Delta t} = i\frac{\hbar}{2m} \frac{1}{2} \left( \frac{u^{(k+1)}_{i+1} - 2u^{(k+1)}_i + u^{(k+1)}_{i-1}}{(\Delta x)^2} + \frac{u^{(k)}_{i+1} - 2u^{(k)}_i + u^{(k)}_{i-1}}{(\Delta x)^2} \right)- \frac{i}{\hbar} V \frac{u^{(k+1)} + u^{(k)}}{2}
# u^{(k+1)}_i - \frac{\alpha}{2} \left( u^{(k+1)}_{i+1} - 2u^{(k+1)}_i + u^{(k+1)}_{i-1} \right) + \frac{\beta}{2} V u^{(k+1)}_i =  u^{(k)}_i + \frac{\alpha}{2} \left( u^{(k)}_{i+1} - 2u^{(k)}_i + u^{(k)}_{i-1} \right) - \frac{\beta}{2} V u^{(k)}_i
# 其中 \alpha = i\frac{\hbar}{2m} \cdot \frac{\Delta t}{(\Delta x)^2}, \quad \beta = \frac{i}{\hbar} \Delta t
# A u^{(k+1)} = B u^{(k)}, \quad
# A = I - \frac{\alpha}{2} \nabla^2 + \frac{\beta}{2} V, \quad
# B = I + \frac{\alpha}{2} \nabla^2 - \frac{\beta}{2} V
# 参考：
# https://wuli.wiki/online/CraNic.html
# https://zhuanlan.zhihu.com/p/393374195
# https://www.bilibili.com/video/BV1sq4y1V7JT
# https://www.bilibili.com/video/BV1CJRvYEEwc
# 可能有bug
# 使用ai辅助
# Gitee Repo

import numpy as np
import scipy
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D

cupy_avaliable = False
gmres_func = None
cupy_dtype = np.complex128
# np.complex128 (2*float64) or np.complex64 (2*float32)
# complex128 can be slow on most consumer gpu, but maybe more accurate

try:
    import cupy
    import cupyx.scipy.sparse
    import cupyx.scipy.sparse.linalg
    cupy_avaliable = cupy.cuda.is_available()
except:
    cupy_avaliable = False;
    
if cupy_avaliable:
    gmres_func = cupyx.scipy.sparse.linalg.gmres
    print('cupy is available')
else:
    gmres_func = scipy.sparse.linalg.gmres
    print('cupy is NOT available')

import time

time0 = time.time()

L = 2;
dx = 0.05;
dt = 0.01;

hbar = 1;
m = 5; # 质量
v = 4;
p = m*v;

x,y,z = np.meshgrid(np.arange(-L,L,dx),np.arange(-L,L,dx),np.arange(-L,L,dx),indexing='ij');
n = int(x.shape[0]);

u0 = np.exp(-20*((x+0.5*L)**2+y**2+z**2))*np.exp(1j*p/hbar*x);
u0[0,:,:]=0;
u0[n-1,:,:]=0;
u0[:,0,:]=0;
u0[:,n-1,:]=0;
u0[:,:,0]=0;
u0[:,:,n-1]=0;

V = np.zeros((n,n,n));
# 双缝
V[np.abs(x)<0.04*L] = 100;
V[(np.abs(y-0.1*L)<0.05*L) | (np.abs(y+0.1*L)<0.05*L)] = 0;

# 单侧高能垒
#V[x>0] = 100;

_u0 = u0.flatten();
_u1 = None;

alpha = 1j*hbar/(2*m)*dt/(dx)**2;
beta = 1j/hbar*dt;

# 构造laplacian算符的矩阵
# 改善效率
def compute_laplacian():
    print('Constructing A...')

    A = scipy.sparse.dok_matrix((n**3,n**3))
    i,j,k = np.meshgrid(np.arange(1,n-1),np.arange(1,n-1),np.arange(1,n-1),indexing='ij')
    i,j,k = i.flatten(),j.flatten(),k.flatten()

    ind = np.zeros(((n-2)**3,7))
    ind[:,0] = np.ravel_multi_index((i,j,k),(n,n,n))
    ind[:,1] = np.ravel_multi_index((i+1,j,k),(n,n,n))
    ind[:,2] = np.ravel_multi_index((i-1,j,k),(n,n,n))
    ind[:,3] = np.ravel_multi_index((i,j+1,k),(n,n,n))
    ind[:,4] = np.ravel_multi_index((i,j-1,k),(n,n,n))
    ind[:,5] = np.ravel_multi_index((i,j,k+1),(n,n,n))
    ind[:,6] = np.ravel_multi_index((i,j,k-1),(n,n,n))
    
    A[ind[:,0],ind[:,0]]=-6
    for l in range(1,7):
        A[ind[:,0],ind[:,l]]=1

    i,j = np.meshgrid(np.arange(0,n),np.arange(0,n),indexing='ij')
    i,j = i.flatten(),j.flatten()
    k0 = np.zeros_like(i)
    k1 = (n-1)*np.ones_like(i)
    ind = np.zeros((n**2,1))
    
    ind[:,0] = np.ravel_multi_index((k0,i,j),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1
    
    ind[:,0] = np.ravel_multi_index((i,k0,j),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1
    
    ind[:,0] = np.ravel_multi_index((i,j,k0),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1
    
    ind[:,0] = np.ravel_multi_index((k1,i,j),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1
    
    ind[:,0] = np.ravel_multi_index((i,k1,j),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1
    
    ind[:,0] = np.ravel_multi_index((i,j,k1),(n,n,n))
    A[ind[:,0],ind[:,0]] = 1

    return A

def compute_Vmat(V):
    print('Constructing Vmat...')
    Vmat = scipy.sparse.dok_matrix((n**3,n**3))
    i,j,k = np.meshgrid(np.arange(0,n),np.arange(0,n),np.arange(0,n),indexing='ij')
    i,j,k = i.flatten(),j.flatten(),k.flatten()
    
    ind = np.ravel_multi_index((i,j,k),(n,n,n))
    Vmat[ind,ind] = np.take(V,ind)

    return Vmat

lap = compute_laplacian()
lap = scipy.sparse.csr_matrix(lap)
Vmat = compute_Vmat(V);

A = scipy.sparse.eye_array(n**3,format = 'csr') - 1/2*alpha*lap + 1/2*beta*Vmat;
B = scipy.sparse.eye_array(n**3,format = 'csr') + 1/2*alpha*lap - 1/2*beta*Vmat;

if cupy_avaliable:
    A = cupyx.scipy.sparse.csr_matrix(A,dtype=cupy_dtype)
    B = cupyx.scipy.sparse.csr_matrix(B,dtype=cupy_dtype)
    _u0 = cupy.asarray(_u0,dtype=cupy_dtype)

def draw(x,y,z,u):
    import plotly.graph_objects as go
    from plotly.offline import plot

    downsample_factor = 3 #降低精度以避免过大的数据

    x_cpu = x[::downsample_factor, ::downsample_factor, ::downsample_factor]
    y_cpu = y[::downsample_factor, ::downsample_factor, ::downsample_factor]
    z_cpu = z[::downsample_factor, ::downsample_factor, ::downsample_factor]
    u_cpu = u[::downsample_factor, ::downsample_factor, ::downsample_factor]

    fig = go.Figure(data=go.Volume(
        x=x_cpu.flatten(),
        y=y_cpu.flatten(),
        z=z_cpu.flatten(),
        value=u_cpu.flatten(),
        opacity=0.1, # needs to be small to see through all surfaces
        surface_count=20, # needs to be a large number for good volume rendering
        ))
    plot(fig)
    return

def draw_slice(x, y, z, u):
    import numpy as np
    import matplotlib.pyplot as plt
    
    x_cpu = x
    y_cpu = y
    z_cpu = z
    u_cpu = u
    
    # 开启交互模式
    plt.ion()
    fig, ax = plt.subplots(figsize=(8, 6))
    
    # 获取中间切片的索引
    n = x_cpu.shape[0]
    slice_idx = int(n / 2)
    
    # 初始化 pcolormesh
    pc = ax.pcolormesh(x_cpu[:, :, slice_idx], 
                       y_cpu[:, :, slice_idx], 
                       u_cpu[:, :, slice_idx], 
                       shading='auto')
    
    plt.xlabel("x") 
    plt.ylabel("y") 
    
    plt.colorbar(pc, ax=ax)
    plt.axis('equal')
    plt.title(f"Slice")
    
    # 立即显示图形
    fig.canvas.draw()
    fig.canvas.flush_events()
    
    return fig, ax, pc  # 返回图形对象以便后续更新
 
def update_slice(u):
    # 转换为 numpy 数组（如果输入是张量）
    u_cpu = u
    
    # 获取中间切片的索引
    n = u_cpu.shape[0]
    slice_idx = int(n / 2)
    
    pc.set_array(u_cpu[:, :, slice_idx].ravel())
    
    fig.canvas.draw()
    fig.canvas.flush_events()

def compute_umag(u):
    return u.real**2+u.imag**2
    
# 绘制初始切片
fig, ax, pc = draw_slice(x, y, z, np.zeros_like(x))

print('Solving u...')
for tick in range(51):
    print(f'tick:{tick}')
    
    _u1 = gmres_func(A,B.dot(_u0))[0]
    _u0,_u1 = _u1,_u0

    if tick % 5 == 0:
        u0 = _u0.get() if cupy_avaliable else _u0;
        u0 = u0.copy()
        u0 = u0.reshape(n,n,n)
        umag = compute_umag(u0)
        update_slice(umag)
        #print(np.sum(umag))
        u0 = None
    

draw(x,y,z,umag)

time1 = time.time()
print(time1-time0)

plt.ioff()
plt.show()
